/*
 * Copyright 2007 Joseph Fifield
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.programmerplanet.ant.taskdefs.jmeter;

import java.io.*;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

import org.apache.tools.ant.BuildException;
import org.apache.tools.ant.DirectoryScanner;
import org.apache.tools.ant.Project;
import org.apache.tools.ant.Task;
import org.apache.tools.ant.taskdefs.Execute;
import org.apache.tools.ant.taskdefs.LogStreamHandler;
import org.apache.tools.ant.types.CommandlineJava;
import org.apache.tools.ant.types.FileSet;

/**
 * Runs one or more JMeter test plans sequentially.
 *
 * @author <a href="mailto:jfifield@programmerplanet.org">Joseph Fifield</a>
 */
public class JMeterConcurrentTask extends Task {

    /**
     * The JMeter installation directory.
     */
    private File jmeterHome;

    /**
     * The property file to use.
     */
    private File jmeterProperties;

    /**
     * Additional user property file to use
     */
    private File additionalUserProperties;

    /**
     * The test plan to execute.
     */
    private File testPlan;

    /**
     * The file to log results to.
     */
    private File resultLog;

    /**
     * The jmeter log file.
     */
    private File jmeterLogFile;

    /**
     * The directory need to save all result log files.
     */
    private File resultLogDir;

    /**
     * A collection of FileSets specifying test plans to execute.
     */
    private ArrayList<FileSet> testPlans = new ArrayList<>();

    private ArrayList<ConcurrentTestPlanSet> concurrentTestPlans = new ArrayList<>();

    /**
     * Whether or not to run the remote servers as specified in the properties file.
     * Default: false.
     */
    private boolean runRemote = false;

    /**
     * The proxy server hostname or ip address.
     */
    private String proxyHost;

    /**
     * The proxy server port.
     */
    private String proxyPort;

    /**
     * The username for the proxy server.
     */
    private String proxyUser;

    /**
     * The password for the proxy server.
     */
    private String proxyPass;

    /**
     * JMeter Log Level
     * */
    private String logLevel;

    /**
     * The main JMeter jar.
     */
    private File jmeterJar;

    /**
     * Array of arguments to be passed to the JVM that will run JMeter.
     */
    private ArrayList<Arg> jvmArgs = new ArrayList<>();

    /**
     * Array of arguments to be passed to JMeter.
     */
    private ArrayList<Arg> jmeterArgs = new ArrayList<>();

    /**
     * Array of properties dynamically passed to JMeter
     */
    private ArrayList<Property> jmProperties = new ArrayList<>();

    /**
     * Indicate if build to be forcefully failed upon testcase failure.
     */
    private String failureProperty;

    /**
     * List of result log files used during run.
     */
    private ArrayList<File> resultLogFiles = new ArrayList<>();

    /**
     * @see org.apache.tools.ant.Task#execute()
     */
    public void execute() throws BuildException {
        this.log("Using " + this.getClass().getName(), 2);

        if (jmeterHome == null || !jmeterHome.isDirectory()) {
            throw new BuildException("You must set jmeterhome to your JMeter install directory.", getLocation());
        }

        jmeterJar = new File(jmeterHome.getAbsolutePath() + File.separator + "bin" + File.separator + "ApacheJMeter.jar");

        validate();

        log("Using JMeter Home: " + jmeterHome.getAbsolutePath(), Project.MSG_VERBOSE);
        log("Using JMeter Jar: " + jmeterJar.getAbsolutePath(), Project.MSG_VERBOSE);

        // 不再支持之前testPlan的配置方式
        if (testPlan != null || testPlans.size() > 0){
            throw new BuildException("Not Support Set Attribute testplan And Element testplans In Class " +
                    this.getClass().getName() + "\nPlease Use Class " + JMeterTask.class.getName(), getLocation());
        }
//        // execute the single test plan if specified
//        if (testPlan != null) {
//            File resultLogFile = resultLog;
//            if (resultLogDir != null) {
//                String testPlanFileName = testPlan.getName();
//                String resultLogFilePath = this.resultLogDir + File.separator + testPlanFileName.replaceFirst("\\.jmx", "\\.jtl");
//                resultLogFile = new File(resultLogFilePath);
//            }
//            executeTestPlan(testPlan, resultLogFile);
//        }

        Iterator<?> testPlanIter;
        // 不再支持之前的testPlans配置方式
//        // execute each of the test plans specified in each of the "testplans" FileSets
//        testPlanIter = testPlans.iterator();
//        while (testPlanIter.hasNext()) {
//            FileSet fileSet = (FileSet)testPlanIter.next();
//            DirectoryScanner scanner = fileSet.getDirectoryScanner(getProject());
//            File baseDir = scanner.getBasedir();
//            String[] files = scanner.getIncludedFiles();
//
//            for (String file : files) {
//                String testPlanFilePath = baseDir + File.separator + file;
//                File testPlanFile = new File(testPlanFilePath);
//                File resultLogFile = resultLog;
//                if (resultLogDir != null) {
//                    String resultLogFilePath = this.resultLogDir + File.separator + file.replaceFirst("\\.jmx", "\\.jtl");
//                    resultLogFile = new File(resultLogFilePath);
//                }
//                executeTestPlan(testPlanFile, resultLogFile);
//            }
//        }

        // 并发执行
        for (ConcurrentTestPlanSet concurrentTestPlanSet : concurrentTestPlans) {
            testPlanIter = concurrentTestPlanSet.getTestPlans().iterator();
            ExecutorService executor = Executors.newCachedThreadPool();
            ArrayList<Future<?>> futures = new ArrayList<>();
            while (testPlanIter.hasNext()) {
                FileSet fileSet = (FileSet) testPlanIter.next();
                DirectoryScanner scanner = fileSet.getDirectoryScanner(getProject());
                File baseDir = scanner.getBasedir();
                String[] files = scanner.getIncludedFiles();

                for (String file : files) {
                    String testPlanFilePath = baseDir + File.separator + file;
                    File testPlanFile = new File(testPlanFilePath);
                    File resultLogFile = resultLog;
                    if (resultLogDir != null) {  // 如果设置了resultlogdir属性，则直接使用jmx文件名作为jtl文件名保存结果
                        String resultLogFilePath = this.resultLogDir + File.separator + file.replaceFirst("\\.jmx", "\\.jtl");
                        resultLogFile = new File(resultLogFilePath);
                    } else {  // 并发执行时不能同时写入同一个jtl文件，按照文件名拆分开
                        String resultLogFilePath = resultLogFile.getParent() + File.separator + file.replaceFirst("\\.jmx", "\\.jtl.tmp");
                        resultLogFile = new File(resultLogFilePath);
                    }
                    Future<?> future = executor.submit(
                            new RunTestPlan(testPlanFile, resultLogFile, jmeterJar, jmeterProperties, jmeterLogFile,
                                    additionalUserProperties, jmeterHome, jvmArgs, jmeterArgs, jmProperties, runRemote,
                                    proxyHost, proxyPort, proxyUser, proxyPass, logLevel, resultLogFiles));
                    futures.add(future);
                }
            }
            for (Future<?> future : futures) {
                try {
                    future.get();  // 等待线程完成
                } catch (InterruptedException | ExecutionException e) {
                    log("Concurrent Execution Failure", Project.MSG_ERR);
                    e.printStackTrace();
                }
            }
            executor.shutdown();
        }

        // 合并jtl文件
        if (resultLogDir == null && resultLog != null){  // 只有满足未设置resultLogDir属性，并设置了resultLog属性才满足合并条件
            try {
                mergeJtl();
            } catch (IOException e) {
                log("Jtl File Merge Failed", Project.MSG_ERR);
                e.printStackTrace();
            }
        }

        checkForFailures();
    }

    /**
    * 合并多个jtl文件
    * */
    private void  mergeJtl() throws IOException {
        final String XML_FIRST_LINE = "<?xml version=\"1.0\" encoding=\"UTF-8\"?>";
        final String XML_SECOND_LINE = "<testResults version=\"1.2\">";
        final String XML_LAST_LINE = "</testResults>";
        final String XML_FIRST_LINE_PATTERN = "<\\?xml version=\"[\\d\\.]+\" encoding=\"[a-zA-Z\\d\\-\\_]+\"\\?>";
        final String XML_SECOND_LINE_PATTERN = "<testResults version=\"[\\d\\.]+\">";
        final String XML_LAST_LINE_PATTERN = "</testResults>";
        final String lineSeparator = System.lineSeparator();

        File resultLogFile = resultLog;

        FileOutputStream fos = new FileOutputStream(resultLogFile);
        OutputStreamWriter osw = new OutputStreamWriter(fos, StandardCharsets.UTF_8);
        BufferedWriter bw = new BufferedWriter(osw);

        bw.write(XML_FIRST_LINE + lineSeparator);
        bw.write(XML_SECOND_LINE + lineSeparator);
        for (File jtlFile : resultLogFiles) {
            log("Merging JTL: " + jtlFile.getAbsolutePath(), Project.MSG_INFO);

            FileInputStream inputStream = new FileInputStream(jtlFile);
            InputStreamReader isr = new InputStreamReader(inputStream, StandardCharsets.UTF_8);
            BufferedReader br = new BufferedReader(isr);
            String line;
            while ((line = br.readLine()) != null) {
                if (isMatch(line, XML_FIRST_LINE_PATTERN) || isMatch(line, XML_SECOND_LINE_PATTERN) || isMatch(line, XML_LAST_LINE_PATTERN)){
                    continue;
                }
                bw.write(line + lineSeparator);
            }
            br.close();
            isr.close();
            inputStream.close();
        }
        bw.write(XML_LAST_LINE + lineSeparator);

        bw.close();
        osw.close();
        fos.close();

        log("All JTL Have Been Merged Into " + resultLogFile.getAbsolutePath(), Project.MSG_INFO);
    }

    /**
     * 判断字符串是否匹配正则
     * */
    private boolean isMatch(String text, String patternStr){
        Pattern pattern = Pattern.compile(patternStr);
        Matcher matcher = pattern.matcher(text);
        return matcher.matches();
    }

    /**
     * Validate the results.
     */
    private void checkForFailures() throws BuildException {
        if (failureProperty != null && failureProperty.trim().length() > 0) {
            for (File resultLogFile : resultLogFiles) {
                log("Checking result log file " + resultLogFile.getName() + ".", Project.MSG_VERBOSE);
                LineNumberReader reader = null;
                try {
                    reader = new LineNumberReader(new FileReader(resultLogFile));
                    // look for any success="false" (pre 2.1) or s="false" (post 2.1)
                    String line = null;
                    while ((line = reader.readLine()) != null) {
                        line = line.toLowerCase();
                        // set failure property if there are failures
                        if (line.indexOf("success=\"false\"") > 0 || line.indexOf(" s=\"false\"") > 0) {
                            log("Failure detected at line: " + reader.getLineNumber(), Project.MSG_VERBOSE);
                            setFailure(getFailureProperty());
                            return;
                        }
                    }
                } catch (IOException e) {
                    throw new BuildException("Could not read jmeter resultLog: " + e.getMessage());
                } finally {
                    try {
                        reader.close();
                    } catch (Exception e) { /* ignore */
                    }
                }
            }
        }
    }

    /**
     * Validate the task attributes.
     */
    protected void validate() throws BuildException {
        if (!(jmeterJar.exists() && jmeterJar.isFile())) {
            throw new BuildException("jmeter jar file not found or not a valid file: " + jmeterJar.getAbsolutePath(), getLocation());
        }

        if (resultLog == null && resultLogDir == null) {
            throw new BuildException("You must set resultLog or resultLogDir.", getLocation());
        }

        if (resultLogDir != null && !(resultLogDir.exists() && resultLogDir.isDirectory())) {
            throw new BuildException("resultLogDir directory not found or not a valid directory: " + resultLogDir.getAbsolutePath(), getLocation());
        }
    }

    /**
     * Execute a JMeter test plan.
     */
    protected void executeTestPlan(File testPlanFile, File resultLogFile) {
        log("Executing test plan: " + testPlanFile + " ==> " + resultLogFile, Project.MSG_INFO);
        resultLogFiles.add(resultLogFile);

        CommandlineJava cmd = new CommandlineJava();

        cmd.setJar(jmeterJar.getAbsolutePath());

        // Set the JVM args
        for (Arg jvmArg : jvmArgs) {
            cmd.createVmArgument().setValue(jvmArg.getValue());
        }

        // Set the JMeter args
        for (Arg jmeterArg : jmeterArgs) {
            cmd.createArgument().setValue(jmeterArg.getValue());
        }

        // non-gui mode
        cmd.createArgument().setValue("-n");
        // the properties file
        if (jmeterProperties != null) {
            cmd.createArgument().setValue("-p");
            cmd.createArgument().setValue(jmeterProperties.getAbsolutePath());
        }
        // the jmeter log file
        if (jmeterLogFile != null) {
            cmd.createArgument().setValue("-j");
            cmd.createArgument().setValue(jmeterLogFile.getAbsolutePath());
        }

        // add user property
        if (additionalUserProperties != null) {
            cmd.createArgument().setValue("-q");
            cmd.createArgument().setValue(additionalUserProperties.getAbsolutePath());
        }
        // the test plan file
        cmd.createArgument().setValue("-t");
        cmd.createArgument().setValue(testPlanFile.getAbsolutePath());
        // the result log file
        cmd.createArgument().setValue("-l");
        cmd.createArgument().setValue(resultLogFile.getAbsolutePath());
        // run remote servers?
        if (runRemote) {
            cmd.createArgument().setValue("-r");
        }

        // the proxy host
        if ((proxyHost != null) && (proxyHost.length() > 0)) {
            cmd.createArgument().setValue("-H");
            cmd.createArgument().setValue(proxyHost);
        }
        // the proxy port
        if ((proxyPort != null) && (proxyPort.length() > 0)) {
            cmd.createArgument().setValue("-P");
            cmd.createArgument().setValue(proxyPort);
        }
        // the proxy user
        if ((proxyUser != null) && (proxyUser.length() > 0)) {
            cmd.createArgument().setValue("-u");
            cmd.createArgument().setValue(proxyUser);
        }
        // the proxy password
        if ((proxyPass != null) && (proxyPass.length() > 0)) {
            cmd.createArgument().setValue("-a");
            cmd.createArgument().setValue(proxyPass);
        }

        // the JMeter runtime properties
        for (Property jmProperty : jmProperties) {
            if (jmProperty.isValid()) {
                cmd.createArgument().setValue((jmProperty.isRemote() ? "-G" : "-J") + jmProperty.toString());
            }
        }

        Execute execute = new Execute(new LogStreamHandler(this, Project.MSG_INFO, Project.MSG_WARN));
        execute.setCommandline(cmd.getCommandline());
        execute.setAntRun(getProject());

        execute.setWorkingDirectory(new File(jmeterHome.getAbsolutePath() + File.separator + "bin"));
        log(cmd.describeCommand(), Project.MSG_VERBOSE);

        try {
            execute.execute();
        }
        catch (IOException e) {
            throw new BuildException("JMeter execution failed.", e, getLocation());
        }
    }

    public void setJmeterHome(File jmeterHome) {
        this.jmeterHome = jmeterHome;
    }

    public File getJmeterHome() {
        return jmeterHome;
    }

    public void setJmeterProperties(File jmeterProperties) {
        this.jmeterProperties = jmeterProperties;
    }

    public File getJmeterProperties() {
        return jmeterProperties;
    }

    public void setTestPlan(File testPlan) {
        this.testPlan = testPlan;
    }

    public File getTestPlan() {
        return testPlan;
    }

    public void setResultLog(File resultLog) {
        this.resultLog = resultLog;
    }

    public File getResultLog() {
        return resultLog;
    }

    public File getJmeterLogFile() {
        return jmeterLogFile;
    }

    public void setJmeterLogFile(File jmeterLogFile) {
        this.jmeterLogFile = jmeterLogFile;
    }

    public void setResultLogDir(File resultLogDir) {
        this.resultLogDir = resultLogDir;
    }

    public File getResultLogDir() {
        return this.resultLogDir;
    }

    public void addTestPlans(FileSet set) {
        testPlans.add(set);
    }

    public void addConcurrentTestPlans(ConcurrentTestPlanSet set) {
        concurrentTestPlans.add(set);
    }

    public ArrayList<ConcurrentTestPlanSet> getConcurrentTestPlans(){
        return concurrentTestPlans;
    }


    public void addJvmarg(Arg arg) {
        jvmArgs.add(arg);
    }

    public void addJmeterarg(Arg arg) {
        jmeterArgs.add(arg);
    }

    public void setRunRemote(boolean runRemote) {
        this.runRemote = runRemote;
    }

    public boolean getRunRemote() {
        return runRemote;
    }

    public void setProxyHost(String proxyHost) {
        this.proxyHost = proxyHost;
    }

    public String getProxyHost() {
        return proxyHost;
    }

    public void setProxyPort(String proxyPort) {
        this.proxyPort = proxyPort;
    }

    public String getProxyPort() {
        return proxyPort;
    }

    public void setProxyUser(String proxyUser) {
        this.proxyUser = proxyUser;
    }

    public String getProxyUser() {
        return proxyUser;
    }

    public void setProxyPass(String proxyPass) {
        this.proxyPass = proxyPass;
    }

    public String getProxyPass() {
        return proxyPass;
    }

    public void setLogLevel(String logLevel) {
        this.logLevel = logLevel;
    }

    public String getLogLevel() {
        return logLevel;
    }

    public void addProperty(Property property) {
        jmProperties.add(property);
    }

    public void setFailureProperty(String failureProperty) {
        this.failureProperty = failureProperty;
    }

    public String getFailureProperty() {
        return failureProperty;
    }

    public void setFailure(String failureProperty) {
        getProject().setProperty(failureProperty, "true");
    }

    public File getAdditionalUserProperties() {
        return additionalUserProperties;
    }

    public void setAdditionalUserProperties(File additionalUserProperties) {
        this.additionalUserProperties = additionalUserProperties;
    }

}
